// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <melon/random.h>
#include <pollux/expression/function_signature.h>
#include <pollux/functions/macros.h>

namespace kumo::pollux::functions {
    template<typename T>
    struct RandFunction {
        static constexpr bool is_deterministic = false;
        static constexpr auto canonical_name = exec::FunctionCanonicalName::kRand;

        MELON_ALWAYS_INLINE void call(double &result) {
            result = melon::Random::randDouble01();
        }

        MELON_ALWAYS_INLINE void call(int64_t &out, const int64_t input) {
            checkBound(input);
            out = melon::Random::rand64(input);
        }

        MELON_ALWAYS_INLINE void call(int32_t &out, const int32_t input) {
            checkBound(input);
            out = melon::Random::rand32(input);
        }

        MELON_ALWAYS_INLINE void call(int16_t &out, const int16_t input) {
            checkBound(input);
            out = int16_t(melon::Random::rand32(input));
        }

        MELON_ALWAYS_INLINE void call(int8_t &out, const int8_t input) {
            checkBound(input);
            out = int8_t(melon::Random::rand32(input));
        }

        template<typename InputType>
        MELON_ALWAYS_INLINE
        typename std::enable_if<std::is_integral<InputType>::value, void>::type
        checkBound(InputType input) {
            POLLUX_USER_CHECK_GT(input, 0, "bound must be positive");
        }
    };

    template<typename T>
    struct SecureRandFunction {
        static constexpr bool is_deterministic = false;

        template<typename InputType>
        MELON_ALWAYS_INLINE void checkInput(
            const InputType lower,
            const InputType upper) {
            POLLUX_USER_CHECK_GT(
                upper, lower, "upper bound must be greater than lower bound");
        }

        MELON_ALWAYS_INLINE void call(double &out) {
            out = melon::Random::secureRandDouble01();
        }

        MELON_ALWAYS_INLINE void
        call(double &out, const double lower, const double upper) {
            checkInput(lower, upper);
            if (std::isinf(lower)) {
                out = std::numeric_limits<double>::quiet_NaN();
            } else if (std::isinf(upper)) {
                out = std::numeric_limits<double>::max();
            } else {
                out = melon::Random::secureRandDouble(0, upper - lower) + lower;
            }
        }

        MELON_ALWAYS_INLINE void
        call(float &out, const float lower, const float upper) {
            checkInput(lower, upper);
            if (std::isinf(lower)) {
                out = std::numeric_limits<float>::quiet_NaN();
            } else if (std::isinf(upper)) {
                out = std::numeric_limits<float>::max();
            } else {
                out = float(melon::Random::secureRandDouble(0, upper - lower)) + lower;
            }
        }

        MELON_ALWAYS_INLINE void
        call(int64_t &out, const int64_t lower, const int64_t upper) {
            checkInput(lower, upper);
            if (upper >= 0) {
                out = int64_t(
                    melon::Random::secureRand64(0, uint64_t(upper) - lower) + lower);
            } else {
                out = int64_t(
                    melon::Random::secureRand64(0, uint64_t(upper - lower)) + lower);
            }
        }

        MELON_ALWAYS_INLINE void
        call(int32_t &out, const int32_t lower, const int32_t upper) {
            checkInput(lower, upper);
            if (upper >= 0) {
                out = int32_t(
                    melon::Random::secureRand32(0, uint32_t(upper) - lower) + lower);
            } else {
                out = int32_t(
                    melon::Random::secureRand32(0, uint32_t(upper - lower)) + lower);
            }
        }

        MELON_ALWAYS_INLINE void
        call(int16_t &out, const int16_t lower, const int16_t upper) {
            checkInput(lower, upper);
            if (upper >= 0) {
                out = int16_t(
                    melon::Random::secureRand32(0, uint16_t(upper) - lower) + lower);
            } else {
                out = int16_t(
                    melon::Random::secureRand32(0, uint16_t(upper - lower)) + lower);
            }
        }

        MELON_ALWAYS_INLINE void
        call(int8_t &out, const int8_t lower, const int8_t upper) {
            checkInput(lower, upper);
            if (upper >= 0) {
                out = int8_t(
                    melon::Random::secureRand32(0, uint8_t(upper) - lower) + lower);
            } else {
                out = int8_t(
                    melon::Random::secureRand32(0, uint8_t(upper - lower)) + lower);
            }
        }
    };
} // namespace kumo::pollux::functions
